import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
import scipy.stats as stats
from tqdm import tqdm
from multiprocessing import Process, Queue
from collections import Counter, defaultdict
from IPython.display import clear_output
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, WhiteKernel

import plotly.io as pio
pio.renderers.default = "notebook+pdf"

Individual Based Generating Processes

Suppose you're trying to figure out how many fish are in a stream. Obviously catching or counting all the fish is impossible. So instead you have to find some indirect way of estimating the number of fish. To do so typically requires coming up with some hypothesis of how your fish are distributed - this is the only way to extrapolate from a subset of your fish to the overall superset. Typically this gets expressed in the following way:

$$P(\vec{s}|\vec{\theta})$$

that is, what is the probability of my sample $\vec{s}$ given my "model parameters" $\vec{\theta}$ where $\vec{\theta}$ usually includes the thing I'm trying to extrapolate (like the total population in the stream in our case). Then using either bayesian methods or maximum likelihood estimation you can figure out the likely $\vec{\theta}$ given this probability distribution and thereby get your estimate.

The key here, obviously is figuring out what $P(\vec{s}|\vec{\theta})$ should be. This we can think of as our "generating process" in the sense that the model underlying $P(\vec{s}|\vec{\theta})$ "generates" our sample $\vec{s}$ with some probability. Now normally these generating processes are rather high level. They might treat the stream as homogenous, or consider an ocean as a big bucket that you sample fish from. They may include things like age structure or mortality, but rarely are the more specific elements of wildlife biology really considered. That is to say that these generating processes are rarely (if ever) grounded in Individual Based Models (IBMs).

Now on the outset this is not necessarily a problem - you can do chemistry without thinking about individual atoms all the time. However in our case things like habitat use, reproduction, the boom and bust of populations given resource limitations, and the like probably have a lot to do with the emergent properties of individual behavior. Of course we could just try to learn these emergent processes themselves but then how do we know how much of that emergence is stable if, say, the environment heats up, or people change their fishes practices, and so on? Furthermore there is already tons of biology on individual organisms and in general, its easier to study a couple of organisms than a whole population so if we could take advantage of individual biology we could probably incorporate and do a lot more science that would then become directly relevant to things like resource management.

So, here is the question, how would one go about creating Individual Based Generating Processes (IBGPs) that could be used for, say, assessment of population size in a stream? Well, let's find out.

Individual Based Models (IBMs)

The very first problem we're going to run into is pretty simple. There are potentially a lot of fish in our stream. And that means if we're going to model each individual we're going to need to deal with scale.

Modeling individuals themselves in not a huge deal because we can always paralellize - i.e. if I have four million fish to deal with I can distribute the compute load across say 100 different machines for a short period of time (things like apache beam allow us to do this pretty easily). However, this all assumes that I can treat my fish independently - yet everything we expect to be interesting is going to come through how fish interact! Therefore not only do I suddenly have a barrier to my paralellization I also have $N^2$ instead of just $N$ things to worry about!

However there's a way we can get around this. And that is by recognizing that at extremely short timescales everything becomes linear. How does this help us? Well it means that we can do things in steps. So at a particular time point we can send the global information to each individual. It can then sort out what of that information is relevant and make changes to its own state on the basis of this. Then it can send its information back to a global pool, and we can repeat. By modeling our nonlinear processes as linear processes at very short timescales we can get our parallelization back.

However, sending the information to each single node can be a bit laborious and we aren't really planning on having a single thread for every single individual so better to group them into units and send and receive information from those units.

So that, in theory will deal with our scale issues. Let's go ahead then and build a very simple IBM!

Let's start by imagining that what we're interested in is how habitat affects the distribution of fish in our stream. In this example we'll assume a couple of things:

1. Fish have a "range" that they patrol that is based on their size.
2. Fish are more likely to capture food from their habitat if they are larger.
3. Food comes in discrete quanta that the fish compete for.
4. Fish need to eat a certain amount of offset their metabolism.
5. Fish migrate toward whatever habitat seems to offer them the best advantages.

From this we can see that we're going to need to model two things - the individual fish, and the individual pieces of habitat (and the food that goes along with them). So let's start with that.

The Fish

class Fish(object):
    def __init__(
        self, tag, position, weight, speed, consumption_rate, unit_territory_radius,
        habitat_centroids, habitat_calories, stream_bounds, time_fraction
    ):
        self.tag = tag
        self.position = position
        self.weight = weight
        self.habitat_centroids = habitat_centroids
        self.habitat_calories = habitat_calories
        self.unit_territory_radius = unit_territory_radius
        self.territory_radius = unit_territory_radius * np.sqrt(self.weight)
        self.speed = speed
        self.reserves = weight
        self.consumption_rate = consumption_rate
        self.stream_bounds = stream_bounds
        self.time_fraction = time_fraction

    def sense(self, state):
        # here we're assuming we've pulled in our
        # global state and that it has the following
        self.fish_weights = np.array(state['weights'])
        self.fish_positions = np.array(state['positions'])
        self.fish_tags = np.array(state['tags'])
        self.won_calories = state['won_calories'].get(self.tag, 0)

    def signal(self):
        # we check to make sure the fish is
        # still alive
        if not self.reserves <= 0:
            # this is the invidual state the fish
            # puts back out into the world
            return 'fish', (self.tag, self.position, self.weight, self.reserves)
        
    def edge_modifier(self, position):
        # compute the actual range of the fish
        # as a proportion of the total allowable
        # range size
        upper = min(position[0] + self.territory_radius, self.stream_bounds[1])
        lower = max(position[0] - self.territory_radius, self.stream_bounds[0])
        return (upper - lower) / (2 * self.territory_radius)
    
    # where all the action happens ;) 
    def act(self):
        # so first we check the metabolic state of
        # the fish, we're going to assume that it's
        # a fraction of the fishes weight. note also
        # that we're multiplying by the time fraction
        # you'll see this all over the place and it 
        # allows us to scale the simulation time
        # (i.e. make our linearity assumption more or
        # less accurate)
        self.reserves -= self.consumption_rate * self.weight * self.time_fraction
        self.reserves += self.won_calories
        if self.reserves <= 0:
            # fish has died so nothing to do 
            return
        
        # alright the next bit of code is literally
        # just going to determine the direction the fish
        # is going to swim in. to do this we need to sort
        # out the competitor fish in each of the surrounding
        # habitats

        # first we need to figure out which fish the current
        # fish can see. we do this by computing the distance
        # to the other fish and the look to see if 
        # the territory of the other fish overlaps with the
        # territory of the current fish
        differences = self.fish_positions - self.position
        distances = np.linalg.norm(differences, axis=1)
        territory_sizes = self.unit_territory_radius * np.sqrt(self.fish_weights)
        interaction_distances = self.territory_radius + territory_sizes

        # filter down to observable fish
        weights = self.fish_weights[(distances <= interaction_distances) & (self.fish_tags != self.tag)]
        positions = self.fish_positions[(distances <= interaction_distances) & (self.fish_tags != self.tag)]
        territory_sizes = self.unit_territory_radius * np.sqrt(weights)

        # next we figure out which habitats this fish
        # is seeing within its range/territory
        differences = self.habitat_centroids - self.position
        habitat_distances = np.linalg.norm(differences, axis=1)
        habitat_centroids = self.habitat_centroids[habitat_distances <= self.territory_radius]
        habitat_calories = self.habitat_calories[habitat_distances <= self.territory_radius]

        # now we can compute the fish's idea of
        # how much each habitat is worth
        best_calories = -float('inf')
        best_centroid = None
        for centroid, calories in zip(habitat_centroids, habitat_calories):
            # figure out which observed fish are in
            # this habitat and get their weights
            differences = positions - centroid
            distances = np.linalg.norm(differences, axis=1)
            sub_weights = weights[distances <= territory_sizes]

            # now we can compute the expected value
            # of being in this habitat. note that we're
            # more or less assuming surrounding habitat 
            # (i.e. that in the fishes range should it move
            # there) is equivalent. therefore we also have to
            # correct for range shrinking when the fish approaches
            # the "edge" of our stream
            expected_calories = self.weight / (np.sum(sub_weights) + self.weight) * calories * self.edge_modifier(centroid)
            if expected_calories > best_calories:
                best_calories = expected_calories
                best_centroid = centroid

        # now that we've got the direction to move towards
        # we can update the position of the fish
        direction = (best_centroid - self.position)
        direction = direction / np.linalg.norm(direction)
        self.position = self.position + direction * self.territory_radius * self.speed * self.time_fraction
        
    

The Habitat

class Habitat(object):
    def __init__(
            self, calories, centroid, unit_territory_radius, calories_per_unit, time_fraction
    ):
        self.calories = np.floor(calories / calories_per_unit) * calories_per_unit
        self.units = int(self.calories / calories_per_unit)
        self.centroid = centroid
        self.unit_territory_radius = unit_territory_radius
        self.calories_per_unit = calories_per_unit
        self.time_fraction = time_fraction
        self.calories_won = {}

    def sense(self, state):
        self.fish_positions = state['positions']
        self.fish_weights = state['weights']
        self.fish_tags = state['tags']

    def signal(self):
        return 'habitat', self.calories_won
    
    def act(self):
        # first we need to determine which fish
        # are competing in this habitat
        differences = self.fish_positions - self.centroid
        distances = np.linalg.norm(differences, axis=1)
        territory_sizes = self.unit_territory_radius * np.sqrt(self.fish_weights)

        # next we compute the probabilities
        # of each fish winning in a competition
        # for calories
        probs = np.ones(len(self.fish_weights))
        # exclude the fish that are too far away
        probs[probs < territory_sizes] = 0
        # bigger fish have a higher chance of winning
        probs = probs * self.fish_weights
        probs = probs / np.sum(probs)

        # now for each unit of calories we're going
        # to randomly assign a winner
        winners = np.random.choice(
            self.fish_tags, size=self.units, p=probs, replace=True
        )

        # assign the spoils
        self.calories_won = {}
        for tag, num_wins in Counter(winners).items():
            self.calories_won[tag] = num_wins * self.calories_per_unit * self.time_fraction

    

Containers and Exchangers

We've now got our individuals (or agents as they'll be called from here on out). Now we need to create the compenents that will actually manage all of these and their states.

class ComputeContainer(object):

    def __init__(self, agents, sub_queue, pub_queue):
        self.pub_queue = pub_queue
        self.sub_queue = sub_queue
        self.agents = agents
        self.read_state = None
        self.write_state = []

    def sense(self):
        self.read_state = self.sub_queue.get()

    def act(self):
        for agent in self.agents:
            agent.sense(self.read_state)
            agent.act()

    def signal(self):
        self.write_state = []
        for agent in self.agents:
            signal = agent.signal()
            if signal:
                self.write_state.append(signal)
        self.read_state = None

        self.pub_queue.put(self.write_state)

class Exchanger(object):
    def __init__(self, sub_queues, pub_queues):
        self.sub_queues = sub_queues
        self.pub_queues = pub_queues
        self.state = None
        self.data = []

    def exchange(self):
        # all our exchanger really does
        # is pull in the states from
        # the compute containers
        # aggregate them and then
        # push them back out to the
        # compute containers
        self.accumulate_state()
        for sub_queue in self.sub_queues:
            sub_queue.put(self.state)
        self.state = None

    def accumulate_state(self):
        self.state = {
            'tags': [],
            'positions': [],
            'weights': [],
            'reserves': [],
            'won_calories': defaultdict(int)
        }
        for pub_queue in self.pub_queues:
            container_states = pub_queue.get()
            for new_state in container_states:
                if new_state[0] == 'habitat':
                    for tag, won_value in new_state[1].items():
                        self.state['won_calories'][tag] += won_value
                else:
                    tag, position, weight, reserves = new_state[1]
                    self.state['tags'].append(tag)
                    self.state['positions'].append(position)
                    self.state['weights'].append(weight)
                    self.state['reserves'].append(reserves)

        self.state['tags'] = np.array(self.state['tags'])
        self.state['positions'] = np.array(self.state['positions'])
        self.state['weights'] = np.array(self.state['weights'])
        self.state['reserves'] = np.array(self.state['reserves'])

        # this here is just a bit of add on
        # so we can grab all the history of our
        # sim after the fact... technically this is
        # not really the job of the exchanger but
        # for simplicity we're just going to do it
        # here instead of making a new agent to
        # accumulate history
        self.data.append(
            pd.DataFrame([
                {
                    'tag': tag,
                    'position': position[0],
                    'weight': weight,
                    'reserves': reserves,
                    'won_calories': self.state['won_calories'].get(tag, 0),
                }
                for tag, position, weight, reserves in zip(
                    self.state['tags'], self.state['positions'], self.state['weights'],
                    self.state['reserves']
                )
            ])
        )

    def retrieve(self):
        # also just a helper function
        # to retrieve the history of
        # the simulation
        for i, df in enumerate(self.data):
            df['time'] = i
        return pd.concat(self.data)

# what our parallel threads will run
def job(agents, trigger_queue, sub_queue, pub_queue):
    container = ComputeContainer(agents, sub_queue, pub_queue)
    container.signal()
    while True:
        trigger = trigger_queue.get()
        if not trigger:
            break
        container.sense()
        container.act()
        container.signal()
    print('process finished')

# how we'll divide up our agents onto threads
def divide_agents(agents, num_processes):
    divided_agents = [list() for _ in range(num_processes)]
    for i, agent in enumerate(agents):
        divided_agents[i % num_processes].append(agent)
    return divided_agents

Putting it All Together

def run_simulation(
    num_agents, consumption_rate, speed, unit_territory_radius, weight_bounds,
    calorie_func, calories_per_unit,
    stream_bounds, habitat_resolution, 
    total_time, time_fraction, 
    num_processes
):
    pub_queues = [Queue() for _ in range(num_processes)]
    sub_queues = [Queue() for _ in range(num_processes)]

    print("initializing habitats")
    habitat_centroids = np.linspace(stream_bounds[0], stream_bounds[1], habitat_resolution)[1:]
    habitat_calories = calorie_func(habitat_centroids, habitat_centroids[1] - habitat_centroids[0])
    habitat_centroids = np.array([[x] for x in habitat_centroids])

    habitats = [
        Habitat(
            calories, centroid, unit_territory_radius, calories_per_unit, time_fraction
        ) for centroid, calories in zip(habitat_centroids, habitat_calories)
    ]
    print("total calories", sum(habitat.calories for habitat in habitats))

    print("created", len(habitats), "habitats")

    print("initializing fish")
    fish = [
        Fish(
            tag, np.array([np.random.uniform(*stream_bounds)]), np.random.uniform(*weight_bounds), 
            speed, consumption_rate, unit_territory_radius, habitat_centroids, habitat_calories,
            stream_bounds, time_fraction
        )
        for tag in range(num_agents)
    ]
    print("consumption per unit time", sum(f.consumption_rate * f.weight for f in fish))

    print('starting processes')
    agents = habitats + fish
    divided_agents = divide_agents(agents, num_processes)
    trigger_queues = [Queue() for _ in range(num_processes)]
    processes = [
        Process(target=job, args=(sub_agents, trigger_queue, sub_queue, pub_queue))
        for sub_agents, trigger_queue, sub_queue, pub_queue in zip(divided_agents, trigger_queues, sub_queues, pub_queues)
    ]
    for process in processes:
        process.start()

    print('building exchanger')
    exchanger = Exchanger(sub_queues, pub_queues)

    print('initializing state')
    exchanger.exchange()

    print('running simulation')
    for _ in tqdm(range(int(round(total_time / time_fraction)))):
        for trigger_queue in trigger_queues:
            trigger_queue.put(1)
        exchanger.exchange()

    print('sending stop signals')
    for trigger_queue in trigger_queues:
        trigger_queue.put(0)

    print('waiting for processes to finish')
    for process in processes:
        process.join()

    print('clearing sub queues')
    while not all(queue.empty() for queue in sub_queues):
        for queue in sub_queues:
            while not queue.empty():
                queue.get()

    print('retrieving results')
    df = exchanger.retrieve()
    return df

Running the Simulation

Let's start with a pretty basic habitat map.

def calories_func(x, radius):
    return x ** 2 * radius * 25

stream_bounds = [0, 1]

df = run_simulation(
    num_agents=100, consumption_rate=0.25, speed=0.25, unit_territory_radius=0.1, weight_bounds=[0.1, 1],
    calorie_func=calories_func, calories_per_unit=0.01,
    stream_bounds=stream_bounds, habitat_resolution=100, 
    total_time=100, time_fraction=1, 
    num_processes=8
)

initializing habitats
total calories 7.980000000000003
created 99 habitats
initializing fish
consumption per unit time 13.28299972298137
starting processes
building exchanger
initializing state
running simulation

100%|██████████| 100/100 [00:00<00:00, 125.57it/s]

process finishedprocess finished
process finished
process finishedprocess finished
process finishedprocess finishedprocess finished

sending stop signals
waiting for processes to finish
clearing sub queues
retrieving results

def watch_simulation(df, color='reserves'):
    jitter = []
    for tag in df['tag'].unique():
        jitter.append({
            'tag': tag,
            'jitter': np.random.uniform(-1, 1)
        })
    jitter = pd.DataFrame(jitter)
    df = df.merge(jitter, on='tag')

    return px.scatter(
        df, x="position", y="jitter", size="weight", color=color, hover_name="tag",
        range_x=[-0.1,1.1], range_y=[-1.1,1.1],
        animation_frame="time", animation_group="tag"
    )

watch_simulation(df)

def watch_props(df, bin_size = 0.1, agg={'weight': 'sum'}):
    df['bin'] = df['position'] // bin_size
    gdf = df.groupby(['bin', 'time']).agg(agg).reset_index().rename({'weight': 'count'}, axis=1)
    full_index = []
    for bin in gdf['bin'].unique():
        for time in gdf['time'].unique():
            full_index.append({
                'bin': bin,
                'time': time
            })
    full_index = pd.DataFrame(full_index)
    gdf = full_index.merge(gdf, on=['bin', 'time'], how='left').fillna(0)
    cbt = gdf.groupby('time').agg({'count': 'sum'}).reset_index().rename({'count': 'total'}, axis=1)
    gdf = gdf.merge(cbt, on='time')
    gdf['proportion'] = gdf['count'] / gdf['total']
    return px.bar(gdf, x='bin', y='proportion', animation_frame='time', range_y=[0, 0.5])

watch_props(df)

Cool! we now have a way of simulating our generating process. Given there's probably a way for us to measure our habitat and we might imagine being able to measure things like metabolic rates we could suppose (once again for simplicity) that our only parameter here really is the number of fish $N$. Furthermore suppose that all we can do is sample the last portion of the stream. The the question is how do we get:

$$P(s|N)$$

We could of course bootstrap it in this case, but that would require a lot of sims across each of the possible $N$ values. However, what if we build a model? Then we could potentially learn a lot from a far smaller number of runs of our simulation. Let's go ahead and give this a try.

Modeling $P(\vec{s}|\vec{\theta})$

Now most models out there just try to predict $\vec{s}$ given $\vec{\theta}$ but we need to actually model the distribution. So what we'll do is take advantage of a special kind of model - Gaussian Process Regression.

Specifically let's start with the simplest model we can think of - a gaussian process regressor based on an RBF kernel and a white kernel (for noise). We'll begin by sampling our simulation a few times.

bin = [0.8, 1.0]
samples = 3
X_train = []
y_train = []
for count in range(10, 61, 5):
    for _ in range(samples):
        df = run_simulation(
            num_agents=count, consumption_rate=0.25, speed=0.25, unit_territory_radius=0.1, weight_bounds=[0.1, 1],
            calorie_func=calories_func, calories_per_unit=0.01,
            stream_bounds=stream_bounds, habitat_resolution=100, 
            total_time=100, time_fraction=1, 
            num_processes=8
        )
        df = df[df['time'] == df['time'].max()]
        df = df[df['position'].between(bin[0], bin[1])]
        X_train.append([count])
        y_train.append(df.shape[0])
        clear_output()

X_train = np.array(X_train)
y_train = np.array(y_train)

px.scatter(x=X_train[:,0], y=y_train, labels={'x': 'total count', 'y': 'collected from bin'})

Now that we have some training data, let's train our Gaussian Process!

best_noise_level = None
best_length_scale = None
best_score = -float('inf')
best_gpr = None
for log_noise_level in tqdm(np.linspace(-3, 3, 10)):
    noise_level = 10 ** log_noise_level
    for log_length_scale in np.linspace(-3, 3, 10):
        length_scale = 10 ** log_length_scale
        white_kernel = WhiteKernel(
            noise_level=noise_level, noise_level_bounds=(1e-5, 1e5)
        )
        rbf_kernel = 1.0 * RBF(length_scale=length_scale, length_scale_bounds=(1e-5, 1e5)) 
        kernel = rbf_kernel + white_kernel
        gpr = GaussianProcessRegressor(kernel=kernel)
        gpr.fit(X_train, y_train)
        score = gpr.log_marginal_likelihood(gpr.kernel_.theta)
        if score > best_score:
            best_score = score
            best_noise_level = noise_level
            best_length_scale = length_scale
            best_gpr = gpr

clear_output()

print('best score', best_score)
print('best noise level', best_noise_level)
print('best_length_scale', best_length_scale)

best score -80.43743926775144
best noise level 1000.0
best_length_scale 10.0

def visualize_fit(X, gpr):
  y_pred, sigma = gpr.predict(X[:, np.newaxis], return_std=True)
  y_upper = y_pred + sigma * 1.96
  y_lower = y_pred - sigma * 1.96

  fig = go.Figure()
  fig.add_trace(go.Scatter(x=X_train[:,0], y=y_train, mode='markers', name='data'))
  fig.add_trace(go.Scatter(x=X, y=y_pred, mode='lines', name='prediction'))
  fig.add_trace(go.Scatter(x=X, y=y_upper, mode='lines', line_color = 'rgba(255, 255, 0, 0.5)'))
  fig.add_trace(go.Scatter(x=X, y=y_lower, mode='lines', name = '95% confidence interval',
                          line_color = 'rgba(255, 255, 0, 0.5)',
                            fill='tonexty', fillcolor = 'rgba(255, 255, 0, 0.2)'))
  return fig

X = np.linspace(1, 65, 100)
visualize_fit(X, best_gpr).show()

Fitting to $s$

Now that we have our generating process as a probability distribution we can invert things:

$$P(N|s) = \frac{P(s|N)P(N)}{P(s)}\propto P(s|N)P(N)$$

Assuming we choose a prior that says we consider all $N$ equally likely this just means that:

$$P(N|s)\propto P(s|N)$$

So suppose we measure 14 fish in that last bin. We can now estimate the expected value and the uncertainty in our measurement.

s = 14
possible_N = np.arange(10, 61, 1)
probabilities = []
for N in tqdm(possible_N):
    X = np.array([N])
    y_pred, sigma = best_gpr.predict(X[:, np.newaxis], return_std=True)
    p = stats.norm.cdf(s + 0.5, y_pred, sigma) - stats.norm.cdf(s - 0.5, y_pred, sigma)
    probabilities.append(p[0])
probabilities = np.array(probabilities)
probabilities = probabilities / np.sum(probabilities)
exp_N = np.sum(possible_N * probabilities)
exp_std = np.sqrt(np.sum((possible_N - exp_N) ** 2 * probabilities))
print(exp_N, exp_std)

100%|██████████| 51/51 [00:00<00:00, 4053.70it/s]

40.97080126110724 11.852761641248335

px.line(x=possible_N, y=probabilities, labels={'x': 'N', 'y': 'likelihood'})

And there you have it! We've been able to generate an estimate of the total population size from our one sample, given a generating function based on a simulation. Obviously we can see this estimate is pretty uncertain but given the data from our simulation this makes sense - that final bin is not particularly well determined by the size of our population.

But the point here is simple, this presents a process for starting from fish behavior and getting all the way to a generating process that can be used for parameter estimation.

And it comes in three steps:

1. Build your simulation.
2. Train a gaussian process regressor for each of the measurements you'll try to make.
3. Use bayesian or MLE methods to flip that regressor on its head and estimate parameters of interest.

One final note is that these simulation generator processes are potentially very powerful because you can generate many different kinds of measurements from them. For example we could look at weights, or number of fish that die off, or other statistics from these bins and then build a GPR for each and thereby take advantage of a lot more information in then determining our parameter estimates. So rather than having a generating process for one $\vec{s}$ you can creating generating processes for many different potential measurements and then munge all of that information together in the end when creating your parameter estimates. Pretty sweet!